#!/usr/bin/env python3
"""
模型路径配置管理器

负责管理ModelScope模型的本地加载路径
支持从本地缓存加载模型，如果本地没有则自动下载到指定目录
"""

import logging
import os
import shutil
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

logger = logging.getLogger(__name__)

@dataclass
class ModelPathConfig:
    """模型路径配置"""
    model_name: str
    model_revision: str
    local_cache_dir: str
    fallback_cache_dir: str
    auto_download: bool = True
    use_local_first: bool = True

class ModelPathManager:
    """
    模型路径管理器

    负责管理ModelScope模型的本地加载路径
    支持从本地缓存加载模型，如果本地没有则自动下载到指定目录
    """

    def __init__(self):
        # 默认缓存目录
        self.default_cache_dir = os.path.expanduser("~/.cache/modelscope/hub/models")
        self.local_cache_dir = os.path.expanduser("~/.cache/modelscope/models")

        # 确保本地缓存目录存在
        os.makedirs(self.local_cache_dir, exist_ok=True)

        # 模型配置映射
        self.model_configs: Dict[str, ModelPathConfig] = {}

        # 初始化模型配置
        self._init_model_configs()

    def _init_model_configs(self):
        """初始化模型配置"""
        from .funasr_config import ModelConfig

        # StructBERT模型配置
        self.model_configs["structbert"] = ModelPathConfig(
            model_name=ModelConfig.ACTION_MODEL,
            model_revision=ModelConfig.ACTION_MODEL_REVISION,
            local_cache_dir=os.path.join(self.local_cache_dir, "iic"),
            fallback_cache_dir=os.path.join(self.default_cache_dir, "iic"),
            auto_download=True,
            use_local_first=True
        )

        # PALM2.0模型配置
        self.model_configs["palm2"] = ModelPathConfig(
            model_name=ModelConfig.SUMMARY_MODEL,
            model_revision=ModelConfig.SUMMARY_MODEL_REVISION,
            local_cache_dir=os.path.join(self.local_cache_dir, "iic"),
            fallback_cache_dir=os.path.join(self.default_cache_dir, "iic"),
            auto_download=True,
            use_local_first=True
        )

        # PALM2.0 Base模型配置
        self.model_configs["palm2_base"] = ModelPathConfig(
            model_name=ModelConfig.TASK_MODEL,
            model_revision=ModelConfig.TASK_MODEL_REVISION,
            local_cache_dir=os.path.join(self.local_cache_dir, "iic"),
            fallback_cache_dir=os.path.join(self.default_cache_dir, "iic"),
            auto_download=True,
            use_local_first=True
        )

    def get_model_path(self, model_type: str) -> Optional[str]:
        """
        获取模型路径

        Args:
            model_type: 模型类型 ("structbert", "palm2", "palm2_base")

        Returns:
            Optional[str]: 模型路径，如果找不到返回None
        """
        if model_type not in self.model_configs:
            logger.error(f"未知的模型类型: {model_type}")
            return None

        config = self.model_configs[model_type]

        # 从模型名称中提取目录名
        model_dir_name = self._extract_model_dir_name(config.model_name)

        # 检查本地缓存目录（支持嵌套目录结构）
        local_path = self._find_model_in_directory(config.local_cache_dir, model_dir_name)
        if config.use_local_first and local_path:
            logger.info(f"使用本地模型: {local_path}")
            return local_path

        # 检查默认缓存目录（支持嵌套目录结构）
        fallback_path = self._find_model_in_directory(config.fallback_cache_dir, model_dir_name)
        if fallback_path:
            logger.info(f"使用默认缓存模型: {fallback_path}")
            # 如果配置了自动复制，则复制到本地目录
            if config.auto_download:
                target_path = os.path.join(config.local_cache_dir, model_dir_name)
                self._copy_model_to_local(fallback_path, target_path)
            return fallback_path

        # 检查临时目录（._____temp）
        temp_path = self._find_model_in_directory(
            os.path.join(self.default_cache_dir, "._____temp", "iic"),
            model_dir_name
        )
        if temp_path:
            logger.info(f"使用临时目录模型: {temp_path}")
            # 复制到本地目录
            if config.auto_download:
                target_path = os.path.join(config.local_cache_dir, model_dir_name)
                self._copy_model_to_local(temp_path, target_path)
            return temp_path

        logger.warning(f"模型 {model_type} 在本地缓存中未找到")
        return None

    def _extract_model_dir_name(self, model_name: str) -> str:
        """
        从模型名称中提取目录名

        Args:
            model_name: 模型名称，如 "iic/nlp_structbert_alimeeting_action-classification_spoken_chinese-base"

        Returns:
            str: 目录名，如 "nlp_structbert_alimeeting_action-classification_spoken_chinese-base"
        """
        return model_name.split("/")[-1]

    def _is_model_available(self, model_path: str) -> bool:
        """
        检查模型是否可用

        Args:
            model_path: 模型路径

        Returns:
            bool: 模型是否可用
        """
        if not os.path.exists(model_path):
            return False

        # 检查是否包含必要的文件
        required_files = ["config.json", "pytorch_model.bin"]
        for file_name in required_files:
            file_path = os.path.join(model_path, file_name)
            if not os.path.exists(file_path):
                return False

        return True

    def _find_model_in_directory(self, base_dir: str, model_dir_name: str) -> Optional[str]:
        """
        在目录中查找模型

        Args:
            base_dir: 基础目录
            model_dir_name: 模型目录名

        Returns:
            Optional[str]: 找到的模型路径，未找到返回None
        """
        if not os.path.exists(base_dir):
            return None

        # 直接查找模型目录
        direct_path = os.path.join(base_dir, model_dir_name)
        if self._is_model_available(direct_path):
            return direct_path

        return None

    def _copy_model_to_local(self, source_path: str, target_path: str) -> bool:
        """
        将模型复制到本地目录

        Args:
            source_path: 源路径
            target_path: 目标路径

        Returns:
            bool: 复制是否成功
        """
        try:
            # 确保目标目录存在
            os.makedirs(os.path.dirname(target_path), exist_ok=True)

            # 如果目标目录已存在，先删除
            if os.path.exists(target_path):
                shutil.rmtree(target_path)

            # 复制模型文件
            shutil.copytree(source_path, target_path)
            logger.info(f"模型已复制到本地: {target_path}")
            return True

        except Exception as e:
            logger.error(f"复制模型失败: {e}")
            return False

    def get_model_config(self, model_type: str) -> Optional[ModelPathConfig]:
        """
        获取模型配置

        Args:
            model_type: 模型类型

        Returns:
            Optional[ModelPathConfig]: 模型配置
        """
        return self.model_configs.get(model_type)

    def list_available_models(self) -> List[Tuple[str, str]]:
        """
        列出可用的模型

        Returns:
            List[Tuple[str, str]]: 模型类型和路径的列表
        """
        available_models: List[Tuple[str, str]] = []
        for model_type, config in self.model_configs.items():
            model_path = self.get_model_path(model_type)
            if model_path:
                available_models.append((model_type, model_path))

        return available_models

    def get_cache_info(self) -> Dict[str, Any]:
        """
        获取缓存信息

        Returns:
            Dict[str, Any]: 缓存信息
        """
        info: Dict[str, Any] = {
            "local_cache_dir": self.local_cache_dir,
            "default_cache_dir": self.default_cache_dir,
            "local_cache_exists": os.path.exists(self.local_cache_dir),
            "default_cache_exists": os.path.exists(self.default_cache_dir),
            "available_models": self.list_available_models(),
            "model_configs": {}  # type: ignore
        }

        for model_type, config in self.model_configs.items():
            model_dir_name = self._extract_model_dir_name(config.model_name)
            local_path = os.path.join(config.local_cache_dir, model_dir_name)
            fallback_path = os.path.join(config.fallback_cache_dir, model_dir_name)

            info["model_configs"][model_type] = {
                "model_name": config.model_name,
                "model_revision": config.model_revision,
                "local_path": local_path,
                "fallback_path": fallback_path,
                "local_available": self._find_model_in_directory(config.local_cache_dir, model_dir_name) is not None,
                "fallback_available": self._find_model_in_directory(config.fallback_cache_dir, model_dir_name) is not None
            }

        return info

# 全局实例
_model_path_manager = None

def get_model_path_manager() -> ModelPathManager:
    """获取全局模型路径管理器实例"""
    global _model_path_manager
    if _model_path_manager is None:
        _model_path_manager = ModelPathManager()
    return _model_path_manager
